package net.prasenjit.test.core.auth;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.PreparedStatementCreatorFactory;
import org.springframework.jdbc.core.RowMapper;
import org.springframework.ldap.core.ContextSource;
import org.springframework.ldap.core.DirContextOperations;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.ldap.userdetails.DefaultLdapAuthoritiesPopulator;
import org.springframework.util.CollectionUtils;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Types;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

public class CustomLdapAuthorizationPopulator extends DefaultLdapAuthoritiesPopulator {

    private static final Logger logger = LogManager.getLogger();

    private String rolePrefix = "ROLE_";
    private String permissionPrefix = "PERMISSION_";
    private boolean convertToUpperCase = true;
    private static final String ROLE_PERMISSION_QUERY = "SELECT RPA.ROLE_NAME, RPA.PERMISSION_NAME "
            + "FROM USER_ROLE_ASSOC URA, ROLE_PERMISSION_ASSOC RPA " + "WHERE URA.ROLE_NAME = RPA.ROLE_NAME "
            + "AND URA.USER_NAME = ?";

    private JdbcTemplate jdbcTemplate;
    private PreparedStatementCreatorFactory preparedStatementCreator;

    public CustomLdapAuthorizationPopulator(ContextSource contextSource, String groupSearchBase) {
        super(contextSource, groupSearchBase);
        preparedStatementCreator = new PreparedStatementCreatorFactory(ROLE_PERMISSION_QUERY, new int[]{Types.VARCHAR});
    }

    @Override
    protected Set<GrantedAuthority> getAdditionalRoles(DirContextOperations user, String username) {
        Set<GrantedAuthority> authorities = new HashSet<GrantedAuthority>();
        try {
            List<RolePermission> rolePermissions = jdbcTemplate.query(
                    preparedStatementCreator.newPreparedStatementCreator(new Object[]{username}),
                    new RowMapper<RolePermission>() {
                        @Override
                        public RolePermission mapRow(ResultSet rs, int rowNum) throws SQLException {
                            RolePermission rolePermission = new RolePermission();
                            rolePermission.setRoleName(rs.getString(1));
                            rolePermission.setPermissionName(rs.getString(2));
                            return rolePermission;
                        }
                    });
            if (!CollectionUtils.isEmpty(rolePermissions)) {
                Set<String> roleSet = new HashSet<String>();
                Set<String> permissionSet = new HashSet<String>();
                for (RolePermission rolePermission : rolePermissions) {
                    GrantedAuthority authority;
                    if (roleSet.add(rolePermission.getRoleName())) {
                        authority = createRoleAuthority(rolePermission.getRoleName());
                        authorities.add(authority);
                    }
                    if (permissionSet.add(rolePermission.getPermissionName())) {
                        authority = createPermissionAuthority(rolePermission.getPermissionName());
                        authorities.add(authority);
                    }
                }
            }
        } catch (DataAccessException e) {
            logger.fatal("Failed to access permissions", e);
        }
        return authorities;
    }

    private GrantedAuthority createRoleAuthority(Object role) {
        if (role instanceof String) {
            if (convertToUpperCase) {
                role = ((String) role).toUpperCase();
            }
            return new SimpleGrantedAuthority(rolePrefix + role);
        }
        return null;
    }

    private GrantedAuthority createPermissionAuthority(Object role) {
        if (role instanceof String) {
            if (convertToUpperCase) {
                role = ((String) role).toUpperCase();
            }
            return new SimpleGrantedAuthority(permissionPrefix + role);
        }
        return null;
    }

    public void setJdbcTemplate(JdbcTemplate jdbcTemplate) {
        this.jdbcTemplate = jdbcTemplate;
    }

    class RolePermission {
        private String roleName;
        private String permissionName;

        public String getPermissionName() {
            return permissionName;
        }

        public void setPermissionName(String permissionName) {
            this.permissionName = permissionName;
        }

        public String getRoleName() {
            return roleName;
        }

        public void setRoleName(String roleName) {
            this.roleName = roleName;
        }
    }

}
